//
// Copyright 2021 Signal Messenger, LLC.
// SPDX-License-Identifier: AGPL-3.0-only
//

use rand::Rng;

#[test]
fn aes_ctr_smoke_test() -> Result<(), signal_crypto::Error> {
    let key = hex::decode("603DEB1015CA71BE2B73AEF0857D77811F352C073B6108D72D9810A30914DFF4")
        .expect("Valid hex");
    let nonce = hex::decode("F0F1F2F3F4F5F6F7F8F9FAFB").expect("Valid hex");
    let init_ctr = 0xFCFDFEFF;
    let input = hex::decode("6BC1BEE22E409F96E93D7E117393172AAE2D8A571E03AC9C9EB76FAC45AF8E5130C81C46A35CE411E5FBC1191A0A52EFF69F2445DF4F9B17AD2B417BE66C3710").expect("Valid hex");
    let output = hex::decode("601EC313775789A5B7A7F504BBF3D228F443E3CA4D62B59ACA84E990CACAF5C52B0930DAA23DE94CE87017BA2D84988DDFC9C58DB67AADA613C2DD08457941A6").expect("Valid hex");

    let mut aes_ctr = signal_crypto::Aes256Ctr32::from_key(&key, &nonce, init_ctr)?;

    let mut buf = input;
    aes_ctr.process(&mut buf)?;

    assert_eq!(hex::encode(buf), hex::encode(output));

    Ok(())
}

#[test]
fn aes_ctr_long_test() -> Result<(), signal_crypto::Error> {
    let key = hex::decode("603DEB1015CA71BE2B73AEF0857D77811F352C073B6108D72D9810A30914DFF4")
        .expect("Valid hex");
    let nonce = hex::decode("FFFFFFFFFFFFFFFFFFFFFFFF").expect("Valid hex");
    let init_ctr = 0x00000000;
    let output = hex::decode("FD4C14729F5004BA49D832AD7BE87C18F4FAFB58962B9A43C3BE41713DED93DBF854AC4CA26285B7F76E04B8F8D4E7D9F7548F9B465C8F713C106E9F63F54305331A4983A2F4B718DE29FA794DA12EEE808642FAEFF8271A0EA28E3CC80EEB65A8EB61F69D8BA97F6BF9054453F55EFB8F9422081F1620FE44ACF99E81122F73D3F921D5E3391654E9947904984375B725FDFBA895C5CDE3D225D7BE3A213C3965178A7DC1E3B552EC7B2FFD9C77EBCC243C4500DFDFBE3B7554AA427C01305BEC48D71AF27C5911D1E649C620D22CF5F3A5AEB9468651DA796F369522FAF91EFABF0FEBD33FCA41C9534606A4EA0199B904B243BA9CB8F37A792DF02EFAB8F0E2E0CF1D579DABA042CFE4C9430AD4EDA786052FCF15E7ACFA2736AAB4590F73675FA1805FE23892C63E0CD01D006935A6E3F8E105A754803D00D9857E49636AB034164156856D58A244EAD475300D93B31E44B5BE3BBF6994EDB895804B4F1BAD43ECFE08B4E130148B669FE620E4F73034FC3E748237870BEC3B1F517684654D1D6BC074DDF7B759A2405F78ED84D1006D25AF9BBC12D6C632F5D543DA0CBE9EA866B2C92126009C27AD59394B76337DE246B50895317E2E345DF3629A5F6227F64522866E7A39121CCC552E3DABC989DCE066DEA355F788C5D92ADA099917A297CFEFA867CE37656FAC6A50798C10B394D5BA54F85CF0F7EF1EEDDFCA1E53E93F1349888CC745190C196F84ECF0721287CC592D406F0A6CC5A55294BF7AA3B35F6CEFC61CAB794B12444312B5E50EC0712E221CC95E9E26E9C3D000881E792AFCB58641B1A94613D64EC72F3DB9AB65BA07A4F05B7E9EE7B335D86A06FCBDB8CBD695AEEF53964A965FFE4C6D7B4E580AB139F8422A702E09EACBEA5D512C31A955B3D60310BE2BBDD73484BAE6612791A19DA3C7B0FD1487E72131A8F9CB801790CE8A6E1E378662CEDCD5EE82BD390576ACFE5334ECD9D907273AEFE67058916388210638E5E60F20EE92389B3533FD6AFFD33095B23D169F0913657F033B8D5C4EA517F167C1D53E031787BBE6D5B577245FFF8151CD8FDCC5D6C32DF70FB8043D42F896CD513B4C85CFF292676CF13B6A1931E87727A561711A3105D9F3519B90C9429B5CD3EDAAE3EE334826A3FD74D6175B5589DB392F956A67C5E67BE59656F1CB37E52C636B2692A60C2044327472FA9AF651AFBCF55D8398A31D343074931A72D54833B29EF21FCB6EF419BB56313513E46C65D833677DBB0F2813E9CE5EF70676102CA0D3C14BBDD659A7498FA08CD359D428A803AEFCC660E9FC704E9BACC5F1D27F2528D46B3FCAA2D47DFA28BF4C").expect("Valid hex");

    let mut aes_ctr = signal_crypto::Aes256Ctr32::from_key(&key, &nonce, init_ctr)?;

    let mut buf = vec![0u8; output.len()];
    aes_ctr.process(&mut buf)?;

    assert_eq!(hex::encode(buf), hex::encode(&output));

    let mut rng = rand::rngs::OsRng;

    for _ in 0..32 {
        // Do it again but with split inputs:
        let mut aes_ctr = signal_crypto::Aes256Ctr32::from_key(&key, &nonce, init_ctr)?;
        let mut buf = vec![0u8; output.len()];

        let mut processed = 0;
        while processed != buf.len() {
            let remaining = buf.len() - processed;
            let this_time = if remaining > 1 {
                rng.gen_range(1, remaining)
            } else {
                remaining
            };
            assert!(this_time > 0);
            aes_ctr.process(&mut buf[processed..processed + this_time])?;
            processed += this_time;
        }

        assert_eq!(hex::encode(buf), hex::encode(&output));
    }

    Ok(())
}
